# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo, leon

import json

from coin.base.timestamp import get_timestamp
from coin.exchange.kraken.kr_rest.public_client import KrakenPublicClient
from coin.exchange.kraken.kr_rest.product import KrakenProduct
from concurrent.futures import ThreadPoolExecutor, as_completed


def try_to_hit_ratelimit_sequencial():
  client = KrakenPublicClient(timeout=10)
  num_query = 0
  num_update = 0
  product = KrakenProduct.FromStr('BCHUSD')
  start_time = get_timestamp()
  duration = 0
  while duration < 120:
    update = client.query_level_book(product)
    num_query += 1
    print(update.res_timestamp - update.req_timestamp)
    print(update.msg)
    if update.msg is None:
      print("status code %s" % update.status)
    elif update.msg and update.status is not None:
      print(json.dumps(update.msg, indent=2))
      break
    else:
      num_update += 1
    duration = (get_timestamp() - start_time) / 1e9
  res = {
      "total_num_query": num_query,
      "total_num_update": num_update,
      "duration": duration,
  }
  return res


def print_stats(stats):
  print("total queries: %d" % stats['total_num_query'])
  print("total updates: %d" % stats['total_num_update'])
  print("duration: %f" % stats['duration'])
  print("average query per second: %f\n" % (stats['total_num_query'] / stats['duration']))


def try_to_hit_ratelimit_parallel(max_workers):
  executor = ThreadPoolExecutor(max_workers=max_workers)
  futures = []
  for _ in range(max_workers):
    futures.append(executor.submit(try_to_hit_ratelimit_sequencial))

  sum_stats = {
      "total_num_query": 0,
      "total_num_update": 0,
      "duration": 0,
  }
  for future in as_completed(futures):
    stats = future.result()
    sum_stats['total_num_query'] += stats['total_num_query']
    sum_stats['total_num_update'] += stats['total_num_update']
    if stats['duration'] > sum_stats['duration']:
      sum_stats['duration'] = stats['duration']
  executor.shutdown()
  print_stats(sum_stats)


if __name__ == '__main__':
  stats = try_to_hit_ratelimit_sequencial()
  print_stats(stats)
  try_to_hit_ratelimit_parallel(20)
